#!/bin/usr/python

import string
import math
import random
import copy

#e constant
e = math.exp(1)
#generate list of characters
chars = list(string.ascii_lowercase)
chars.extend(['.',',','?','!',' '])
#generate map of quadgrams and log probability for fitness scoring
quadgrams = {}
N = 0
floor = 0
with open('english_quadgrams.txt','r') as ifile:
    lines = ifile.readlines()
    N = len(lines)
    floor = math.log10(0.01/N)
    for i in range(0,len(lines)):
        q,f = string.split(lines[i])
        quadgrams[q.lower()] = math.log10(float(f)/N)
    ifile.close()


class crackPolygraphicCipher:

    def __init__(self,n):
        #size of block
        self.n = n

        #initial key: randomized
        keys = []
        for i in range(0,len(chars)):
            for j in range(0,len(chars)):
                for k in range(0,len(chars)):
                    for l in range(0,len(chars)):
                        s = [chars[i],chars[j],chars[k],chars[l]]
                        keys.append(''.join(s))
        vals = list(keys)
        random.shuffle(vals)
        self.key = dict(zip(keys,vals))

    #attempt to crack a string 'message'
    #does so by attempting ciphers, checking the fitness of the decrypted
    #message, changing the cipher based off of that fitness, and trying again
    #variation of simulatted annealing
    def crack_string(self,message):
        #initial temperature, fitness, and test key (ckey)
        t = 10
        fitness = -100000000000
        ckey = copy.deepcopy(self.key)

        #create the list of n-grams in the message
        #these are the only n-grams we will attempt to change in ckey
        #because nothing else will affect the decrypted message
        #'vals' holds either 1 or 0. this is used later to determine
        #whether we want to change the key for that gram. if the gram
        #has a good fitness score, we do not want to change it. initilized
        #to '0' (change gram) for all grams
        j = 0
        grams = []
        vals = []
        while j < len(message):
            m = ''
            if j+self.n >= len(message):
                if j == len(message)-1:
                    m = message[j]
                else:
                    m = message[j:len(message)]
                i = self.n+j-len(message)
                while i > 0:
                    i -= 1
                    m += ' '
            else:
                m = message[j:j+self.n]
            grams.append(m)
            vals.append(0)
            j += self.n
        self.grams = dict(zip(grams,vals))

        #copy of keys for easy access
        keys = self.key.keys()

        #simulated annealing loop
        #while loop is used instead of a for loop for flexibility
        #so it is possible to adjust the step by a value other than 1
        while t > 0:
            for i in range(0,50000):
                #change the values in ckey to try another decryption
                #if we decide the change self.key because we have a
                #higher fitness score, we do not want to perform a 
                #deep copy over the entire key because that takes too long.
                #instead we keep track of the values we have changed in nVals
                nVals = []
                for j in self.grams.keys():
                    #only change key if the gram does not already
                    #have a high score
                    if self.grams[j] == 0:
                        x = random.randint(0,923520)
                        n = keys[x]
                        nVals.append(n)
                        tmp = ckey.get(j)
                        ckey[j] = ckey.get(n)
                        ckey[n] = tmp

                #decrypt using potential key 'ckey'
                m = self.decrypt(ckey,message)
                #check fitness score of new message
                mF = self.determineFitness(m,message)
#                print str(i) + ': ' + str(mF) + ' ' + str(fitness)
#                print m
                #if we've found a better fitness, update the key
                if mF - fitness > 0:
                    fitness = mF
                    #update self.key
                    for j in range(0,len(nVals)):
                        x = self.grams.keys()[j]
                        if self.grams[x] == 0:
                            self.key[x] = copy.deepcopy(ckey.get(x))
                            self.key[nVals[j]] = copy.deepcopy(ckey.get(nVals[j]))
                #if we haven't, take the new fitness / key with a probability p
                elif mF - fitness < 0:
                    change = (mF - fitness) / t
                    p = math.pow(e,change)
                    if random.random() < p:
                        fitness = mF
                        #update self.key
                        for j in range(0,len(nVals)):
                            x = self.grams.keys()[j]
                            if self.grams[x] == 0:
                                self.key[x] = copy.deepcopy(ckey.get(x))
                                self.key[nVals[j]] = copy.deepcopy(ckey.get(nVals[j]))
            #decrement temperature
            t -= 1

        return self.decrypt(self.key,message)
        
    #decrypt 'message' using 'key'. breaks up the message into n-grams
    #and finds the corresponding value in key
    def decrypt(self,key,message):
        m = ''
        for i in self.grams.keys():
            d = key.get(i)
            m += ''.join(d)
        #join string back together
        return m

    #determines the fitness of a statement
    #breaks message into overlapping quadgrams and compares the
    #frequency of those quadgrams to a list of frequencies in
    #war and peace.
    #we pass the original message so we can mark grams as found
    def determineFitness(self,message,omessage):
        #strip message of punctation
        message = message.replace(' ','')
        message = message.replace('.','')
        message = message.replace(',','')
        message = message.replace('!','')
        message = message.replace('?','')
        s = 0
        for i in range(0,len(message)-3):
            if message[i:i+4] in quadgrams:
                s += quadgrams[message[i:i+4]]
                #if score is below a threshold, then we no 
                #longer want to change the value
                if quadgrams[message[i:i+4]] > -1:
                    x = int(math.floor(i/4) * 4)
                    self.grams[omessage[x:x+4]] = 1
                #if it is above the threshold, we want to continue testing
            else:
                s += floor
        return s
